package io.yanglong.aiassistant;

import dev.langchain4j.data.message.AiMessage;
import dev.langchain4j.data.message.UserMessage;
import dev.langchain4j.model.chat.ChatLanguageModel;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.input.PromptTemplate;
import dev.langchain4j.model.ollama.OllamaChatModel;
import dev.langchain4j.model.ollama.OllamaEmbeddingModel;
import dev.langchain4j.model.output.Response;
import dev.langchain4j.rag.content.Content;
import dev.langchain4j.rag.content.injector.ContentInjector;
import dev.langchain4j.rag.content.injector.DefaultContentInjector;
import dev.langchain4j.rag.content.retriever.ContentRetriever;
import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
import dev.langchain4j.rag.query.Query;
import dev.langchain4j.store.embedding.redis.RedisEmbeddingStore;
import io.yanglong.aiassistant.config.Constant;

import java.util.List;

public class MeiTuanRag {
    public static void main(String[] args) {
        String userMsg = "在线支付取消订单后钱怎么返还？";
        //1.将问题向量化，然后去向量数据库中查询相关度比较高的知识条目
        EmbeddingModel embeddingModel = OllamaEmbeddingModel.builder()
                .baseUrl("http://127.0.0.1:11434")
                .modelName(Constant.MODEL_NAME_QW2_5_7B)
                .build();
        RedisEmbeddingStore embeddingStore = RedisEmbeddingStore.builder()
                .host("127.0.0.1")
                .port(6379)
                //维度，根据模型设置
                .dimension(Constant.DS_7B_QWEN_DIMENSION)
                .indexName("meituan-rag")
                .build();
        ContentRetriever contentRetriever = EmbeddingStoreContentRetriever.builder()
                .embeddingModel(embeddingModel)
                .embeddingStore(embeddingStore)
                //最多返回5条知识条目
                .maxResults(5)
                //返回相关评分高于此值的知识条目，只支持0.0到1.0
                //.minScore(0.1d)
                .build();
        List<Content> retrieved = contentRetriever.retrieve(Query.from(userMsg));
        for (Content content : retrieved) {
            System.out.println(content);
        }
        //2.将知识条目和问题整合成一个完整的信息，传递给大模型
        ContentInjector contentInjector = DefaultContentInjector.builder()
                //提示此模板，这个非常重要，可以使用变量：{{userMessage}}表示用户输入的信息，{{contents}}表示从向量数据库中查出来的知识条目。不同的模板会影响大模型生成的效果
                .promptTemplate(PromptTemplate.from("""
                        你是一个中文问答机器人。
                        你的任务是根据下面给定的已知信息回答用户的问题。
                        已知信息：
                        {{contents}}\n
                        用户的问题：
                        {{userMessage}}
                        
                        如果已知信息不包含用户问题的答案，或者已知信息不足以回答用户的问，请直接回复‘我无法回答您的问题’。
                        请不要输出已知信息中不包含的信息或答案。
                        """))
                .build();
        UserMessage promptMsg = contentInjector.inject(retrieved, UserMessage.from(userMsg));
        System.out.println(promptMsg.singleText());
        System.out.println("==============================================================");
        ChatLanguageModel chatLanguageModel= OllamaChatModel.builder()
                .baseUrl("http://127.0.0.1:11434")
                .modelName(Constant.MODEL_NAME_QW2_5_7B)
                .maxRetries(3)
                .build();
        Response<AiMessage> aiMessageResponse = chatLanguageModel.generate(promptMsg);
        System.out.println(aiMessageResponse.content().text());
    }
}
